import os
import chromadb
from langchain_community.embeddings import QianfanEmbeddingsEndpoint

os.environ["QIANFAN_AK"] = "SGbbQdjFjlKurTfUIjYM0Q4P"
os.environ["QIANFAN_SK"] = "lb1tKvDGRhqLZYH4ZYpke6Vco9n9X8Xv"

class QianfanEmbeddingAdapter:
    def __init__(self):
        self.client = QianfanEmbeddingsEndpoint()

    def __call__(self, input: list[str]) -> list[list[float]]:
        # 确保 input 是 list[str]
        if isinstance(input, str):
            input = [input]
        # Qianfan 的 __call__ 支持 *args，改用 client.__call__，传入 input 参数列表
        embeddings = self.client(input)  # 这里直接调用 __call__
        return embeddings

embed_model = QianfanEmbeddingAdapter()

client = chromadb.Client()
collection = client.create_collection(
    name="my_knowledge_base",
    metadata={"hnsw:space": "cosine"},
    embedding_function=embed_model
)

collection.add(
    documents=["RAG是一种检索增强生成技术", "向量数据库存储文档的嵌入表示", "三英战吕布"],
    metadatas=[{"source": "tech_doc"}, {"source": "tutorial"}, {"source": "tutorial1"}],
    ids=["doc1", "doc2", "doc3"]
)

results = collection.query(
    query_texts=["什么是RAG技术？"],
    n_results=3
)
print("查询结果：", results)
